Contents

The Maxwell-Boltzmann distribution

This example is designed to stress-test heyoka.py’s event detection system by simulating a two-dimensional box containing ideal gas particles undergoing fully-elastic collisions. Statistical mechanics tells us that, regardless of the initial conditions, the velocity distribution of the gas particles will eventually converge to a 2D Maxwell-Boltzmann distribution.

Before starting, we need to emphasize that heyoka.py is not the best tool for this type of simulation: the dynamics is trivial (rectilinear uniform motion) and collision detection/handling for a large number of particles is more efficiently performed with the help of, e.g., a quadtree or spatial hashing. Nevertheless, this example will show how heyoka.py is able to deal with a large number of events, and it will also illustrate a few caveats regarding the integration of systems with trivial dynamics.

Setting things up

We begin by introducing a few constants for our simulation:

# Square root of the number of particles.
N_sqrt = 25

# Number of particles.
N = N_sqrt * N_sqrt

# Particle radius.
p_radius = .75

# Box size.
box_size = N_sqrt * 10.

For this simulation, we will use \(N=25\times 25 = 625\) particles. The ODEs governing the particles’ motion are trivial:

\[\begin{split} \begin{cases} \frac{dx_i}{dt} = v_{x,i}\\ \frac{dy_i}{dt} = v_{y,i}\\ \frac{dv_{x,i}}{dt} = 0\\ \frac{dv_{y,i}}{dt} = 0 \end{cases}, \end{split}\]

with \(i \in \left[0, N\right)\). Let us begin by implementing these ODEs:

import heyoka as hy

# The list of ODEs.
eqns = []

for i in range(N):
    xi, yi, vxi, vyi = hy.make_vars("x_{}".format(i), "y_{}".format(i), "vx_{}".format(i), "vy_{}".format(i))
    eqns.append((xi, vxi))
    eqns.append((yi, vyi))
    eqns.append((vxi, hy.expression(0.)))
    eqns.append((vyi, hy.expression(0.)))

In order to handle particle collisions, we will have to define \(N\times\left( N - 1 \right) / 2\) events which will:

  • detect when two particles collide,

  • react to the collision by flipping the velocity components across the surface of contact.

For each pair of particles \(\left( i, j \right), i\neq j\) the sphere-sphere collision event equation is:

\[ \left(x_i - x_j\right)^2 + \left(y_i - y_j\right)^2 - 4\mathrm{p\_radius}^2 = 0. \]

We will also have to handle the particles’ collisions with the box, whose event equations are:

\[\begin{split} \begin{cases} x_i \pm \left(\frac{\mathrm{box\_size}}{2} - \mathrm{p\_radius}\right) = 0 \\ y_i \pm \left(\frac{\mathrm{box\_size}}{2} - \mathrm{p\_radius}\right) = 0 \end{cases}. \end{split}\]

When a particle hits a wall, the corresponding velocity component will be flipped.

Let us see the implementation of the callbacks first:

import numpy as np

# Collisions with the left/right walls.
class cb_left_right:
    def __init__(self, idx):
        # Store the particle index on construction.
        self.idx = idx

    def __call__(self, ta, mr):
        # Reshape for ease of indexing.
        st = ta.state.reshape((N, 4))
        
        # Flip the x component of the velocity
        # for the affected particle.
        st[self.idx, 2] = -st[self.idx, 2]

        # Return True to continue the integration.
        return True

# Collisions with the top/bottom walls.
class cb_top_bottom:
    def __init__(self, idx):
        # Store the particle index on construction
        self.idx = idx

    def __call__(self, ta, mr):
        # Reshape for ease of indexing.
        st = ta.state.reshape((N, 4))

        # Flip the y component of the velocity
        # for the affected particle.
        st[self.idx, 3] = -st[self.idx, 3]

        # Return True to continue the integration.
        return True

# Sphere-sphere collision.
class cb_sph_sph:
    def __init__(self, i, j):
        # Store the indices of the particle pair.
        self.i = i
        self.j = j

    def __call__(self, ta, mr):
        # Reshape for ease of indexing.
        st = ta.state.reshape((N, 4))
        
        # Determine the unit vector uij
        # connecting particle i to particle j.
        ri = st[self.i, 0:2]
        rj = st[self.j, 0:2]
        rij = rj - ri
        uij = rij / np.linalg.norm(rij)

        # Fetch the velocities of the
        # two particles.
        vi = st[self.i, 2:4]
        vj = st[self.j, 2:4]
        
        # Porject vi/vj across uij.
        proj_i = np.dot(vi, uij)
        proj_j = np.dot(vj, uij)
        
        # Flip the velocity components
        # across uij.
        vi -= proj_i * uij
        vj -= proj_j * uij
        vi += proj_j * uij
        vj += proj_i * uij
        
        # Return True to continue the integration.
        return True

Next, we will define the event equations and create the events:

# The list of events.
t_events = []

# Collisions with the box walls.
t_events += [hy.t_event(hy.expression("y_{}".format(i)) + (box_size/2 - p_radius), callback = cb_top_bottom(i), direction=hy.event_direction.negative) for i in range(N)]
t_events += [hy.t_event(hy.expression("y_{}".format(i)) - (box_size/2 - p_radius), callback = cb_top_bottom(i), direction=hy.event_direction.positive) for i in range(N)]
t_events += [hy.t_event(hy.expression("x_{}".format(i)) + (box_size/2 - p_radius), callback = cb_left_right(i), direction=hy.event_direction.negative) for i in range(N)]
t_events += [hy.t_event(hy.expression("x_{}".format(i)) - (box_size/2 - p_radius), callback = cb_left_right(i), direction=hy.event_direction.positive) for i in range(N)]

# Sphere-sphere collisions.
for i in range(N):
    xi, yi = hy.make_vars("x_{}".format(i), "y_{}".format(i))
    for j in range(i + 1, N):
        xj, yj = hy.make_vars("x_{}".format(j), "y_{}".format(j))
        
        t_ev = hy.t_event((xi - xj)**2 + (yi - yj)**2 - 4*p_radius**2, callback = cb_sph_sph(i, j), direction = hy.event_direction.negative)
        
        t_events.append(t_ev)

Note how, as explained in the Keplerian billiard example, we assigned specific directions to the collision events. This is done in order to avoid triggering spurious double-bounce events that would lead to the particles exiting the box or penetrating one another.

We are now ready to create the integrator. We will set all positions and velocities initially to zero and enable compact_mode - this is important, otherwise the creation of the integrator object for such a large system would take an excessive amount of time. We also set the integrator tolerance to a high value, because the dynamics is trivial and using a small tolerance would only slow things down without providing any benefit:

ta = hy.taylor_adaptive(eqns, np.zeros(shape=(N * 4)), compact_mode = True, t_events = t_events, tol = 1e-2)

We can now set up the particles’ initial conditions. In order to avoid overlaps in the initial positions, the particles will be laid out in a regular grid, and their velocities will all have a magnitude of 1 and a randomly-generated direction:

# Reshape for ease of indexing.
st = ta.state.reshape((N, 4))

# Randomly-generate unitary velocity vectors.
v_arr = np.random.default_rng().random((N,2)) * 2 - 1
v_arr = v_arr / np.repeat(np.linalg.norm(v_arr, axis=1), 2).reshape((N, 2))
st[:, 2:4] = v_arr

# Place the initial positions in a grid.
pos_arr = np.array([(i, j) for i in range(N_sqrt) for j in range(N_sqrt)])
st[:, 0:2] = (pos_arr * 10) + 5 - box_size / 2

Let us also define a grid of times for the numerical integration:

Ngrid = 200
t_grid = np.linspace(0, Ngrid, Ngrid)

We are now ready to begin the numerical integration.

Running the simulation

For this specific example, it is important to clamp the integration timestep via the max_delta_t keyword argument: because the dynamics is trivial, heyoka.py would attempt to take very long timesteps which would in turn greatly slow down collision detection. Additionally, we will also use tqdm in conjunction with the integrator’s callback mechanism to provide a fancy progressbar that will inform us about the status of the integration:

from tqdm.auto import tqdm

# Progressbar callback.
class cb:
    def __init__(self, tot):
        self.pbar = tqdm(total = tot, unit_scale=True)
        self.cur = 0.
        self.tot = tot
    def __call__(self, ta):
        dt = ta.time - self.cur
        self.pbar.update(ta.time - self.cur)
        self.cur += dt

# Run the integration over the time grid.
oc, _, _, _, res = ta.propagate_grid(t_grid, max_delta_t = 5, callback = cb(float(Ngrid)))

We can now take a look at the simulation results, and, in particular, at the evolution in time of the velocity distribution:

from matplotlib.pylab import plt
from matplotlib import animation, rc
plt.rcParams["figure.figsize"] = (24,12)
plt.rcParams["animation.embed_limit"] = 40
from IPython.display import HTML

fig, axes = plt.subplots(1, 2)
ax1, ax2 = axes

c_list = []
for i in range(N):
    cc = plt.Circle((st[i, 0], st[i, 1]), p_radius, ec='black', fc='orange', zorder=2)
    ax1.add_artist(cc)
    c_list.append(cc)

def init():
    ax1.add_patch(plt.Rectangle((-box_size/2, -box_size/2), box_size, box_size, edgecolor='k', facecolor="none", linestyle = '--'))
    ax1.axis('equal')
    ax1.set_xlim((-box_size/2 - box_size/2/10.,box_size/2 + box_size/2/10.))
    ax1.set_ylim((-box_size/2 - box_size/2/10.,box_size/2 + box_size/2/10.))   
    return tuple(c_list)

def animate(i):  
    fig.suptitle("$t=${:.0f}".format(t_grid[i]))
    st = res[i].reshape((N, 4))
    
    for j in range(N):
        c_list[j].set_center((st[j, 0], st[j, 1]))

    ax2.clear()
    _, bins, _ = ax2.hist(np.linalg.norm(st[:, 2:4], axis=1), range=(0, 3), bins=20, alpha=.6, density=True)
    xgr = np.linspace(0, bins[-1], 200)
    ax2.plot(xgr, 2*xgr*np.exp(-xgr*xgr), 'k--')
    ax2.set_yticks([])
    ax2.set_ylim(0, 1.2)

    return tuple(c_list)

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=Ngrid, interval=50, 
                               blit=True)

%matplotlib inline

HTML(anim.to_jshtml())

As the animation shows, the velocity distribution quickly settles into the Maxwell-Boltzmann distribution (dashed line in the right panel) as a result of the particles’ collisions.

The wavy ramp Examples